package com.hbny.dynamicDataSource.config;

import com.hbny.dynamicDataSource.exception.ApplicationException;
import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DriverManager;
import java.util.HashMap;
import java.util.Map;

/**
 * @author : Aaron
 * @date : 2022/6/2
 * （切换数据源必须在调用service之前进行，也就是开启事务之前）
 */
public class DynamicDataSource extends AbstractRoutingDataSource {

    private static final Logger log = LoggerFactory.getLogger(DynamicDataSource.class);

    private static Map<Object, Object> targetDataSources = new HashMap<>();


    /**
     * 如果不希望数据源在启动配置时就加载好，可以定制这个方法，从任何你希望的地方读取并返回数据源
     * 比如从数据库、文件、外部接口等读取数据源信息，并最终返回一个DataSource实现类对象即可
     */
    @Override
    protected DataSource determineTargetDataSource() {
        return super.determineTargetDataSource();
    }

    /**
     * 如果希望所有数据源在启动配置时就加载好，这里通过设置数据源Key值来切换数据，定制这个方法
     */
    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.getDataSourceKey();
    }

    /**
     * 设置默认数据源
     *
     * @param defaultDataSource
     */
    @Override
    public void setDefaultTargetDataSource(Object defaultDataSource) {
        super.setDefaultTargetDataSource(defaultDataSource);
    }

    /**
     * 设置数据源
     *
     * @param dataSources
     */
    @Override
    public void setTargetDataSources(Map<Object, Object> dataSources) {
        super.setTargetDataSources(dataSources);
        this.targetDataSources = dataSources; //为了可以动态添加数据源
        // 将数据源的 key 放到数据源上下文的 key 集合中，用于切换时判断数据源是否存在
        DynamicDataSourceContextHolder.addDataSourceKeys(dataSources.keySet());
    }

    /**
     * 动态增加数据源
     *
     * @param dbConfig 数据源属性
     * @param key      租户唯一标识
     * @return
     */
    public synchronized boolean addDataSource(String key, HikariConfig dbConfig) {
        if (StringUtils.isBlank(key)) {
            return false;
        }
        try {
            //初始化连接一次，排除配置数据库异常
            Connection connection = null;
            try {
                Class.forName(dbConfig.getDriverClassName());
                connection = DriverManager.getConnection(dbConfig.getJdbcUrl(), dbConfig.getUsername(), dbConfig.getPassword());
                log.info("----database connect test----{}，{}", key, connection.isClosed());
            } catch (Exception e) {
                throw new ApplicationException(5000, "数据库配置异常");
            } finally {
                if (connection != null && !connection.isClosed()) {
                    connection.close();
                }
            }
            if (DynamicDataSource.isExistDataSource(key)) {
                return true;
            }
            HikariDataSource dataSource = new HikariDataSource(dbConfig);
            targetDataSources.put(key, dataSource);
            afterPropertiesSet();
            log.info("----dataSource {} has been added----", key);
        } catch (Exception e) {
            log.error(e.getMessage());
            return false;
        }
        return true;
    }


    /**
     * 是否存在当前key的 DataSource
     *
     * @param key
     * @return 存在返回 true, 不存在返回 false
     */
    public static boolean isExistDataSource(String key) {
        return targetDataSources.containsKey(key);
    }

}